import os
import time

import torchvision
import torch
import cv2
import numpy as np

import evaluation
# from model import MyFCN
from modelLRASSP import MyLRASSP
# from modelDeeplab import MyDeeplab


if __name__ == '__main__':
    # 设备选定cuda
    device = torch.device('cuda')
    # 数据转换，与数据集的转换要求一致
    transformer = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),  # 张量
        # torchvision.transforms.Resize((1250, 500)), # tochvivison格式的size为（H，W）与opencv（W，H）相反（测试模型，单张图片不要求大小）
        torchvision.transforms.Normalize(mean=[0.734, 0.734, 0.7340], std=[0.0917, 0.0917, 0.0917])  # 归一化
    ])
    # 创建模型结构实例
    mylrassp = MyLRASSP()
    # mylrassp = MyFCN()

    # 导入模型参数
    mylrassp.load_state_dict(torch.load(r'./lr_assp_best.pt'))
    mylrassp.to(device)
    # 测试模型状态
    mylrassp.eval()

    # 测试数据
    # img_size = (500, 1250)  # opencv格式的size（W，H）（测试模型单张图片不要求大小）
    img_path = r'./data/test/imgs'
    mask_path = r'./data/test/masks'
    img_names = os.listdir(img_path)
    mask_names = os.listdir(mask_path)

    # 创建相关窗口
    cv2.namedWindow('mask real', cv2.WINDOW_NORMAL)
    cv2.namedWindow('mask', cv2.WINDOW_NORMAL)
    cv2.namedWindow('img', cv2.WINDOW_NORMAL)

    miou_sum = 0
    # 依次读取每张图片进行测试
    for file in img_names:
        # 读取图片
        img = cv2.imread(os.path.join(img_path,file))
        print('img.shape: ', img.shape)
        img_copy = img.copy()

        # 读取标签，注意标签是灰度图，黑色0，白色255
        mask_real = cv2.imread(os.path.join(r'./data/test/masks', os.path.splitext(file)[0]+'.bmp'), 0)
        mask_real_copy = mask_real.copy()
        # 把真实标签像素设为id像素，以配合模型输出计算评价指标
        mask_real[mask_real<=128] = 0
        mask_real[mask_real>128] = 1
        mask_real_tensor = torch.Tensor(mask_real).reshape(-1, mask_real.shape[1], mask_real.shape[0])  # 转为tensor方便移动设备

        # 输入数据预处理
        img = transformer(img)
        img = img.reshape(-1,img.shape[0], img.shape[1], img.shape[2])  # 注意输入模型含有batchsize
        img = img.to(device)

        # 模型前向传播
        with torch.no_grad():
            t = time.time()
            result = mylrassp(img)['out']  # 注意语义分割取'out'
            print('forward time: ', 1000*(time.time() - t))


        # 取结果
        result = result.argmax(dim = 1)  # （batch_size, H, W）

        # 移动至CPU以计算评价指标
        result = result.to('cpu')

        # 计算评价指标
        CM = evaluation.confusion_matrix(mask_real_tensor, result)
        acc = evaluation.pixelAccuracy(CM)
        clsacc = evaluation.clsPixelAccuracy(CM)
        miou = evaluation.MIOU(CM)
        miou_sum += miou
        print('CM: ', CM, sep = '\n')
        print('acc: ', acc)
        print('clsacc: ', clsacc)
        print('miou: ', miou)
        print('-----------------------------------------')

        # 计算结果
        mask = np.zeros(shape=img_copy.shape[:2], dtype=np.uint8)
        mask[result[0,:,:]==1] = 255

        # 显示图片
        cv2.imshow('img', img_copy)
        cv2.imshow('mask', mask)
        cv2.imshow('mask real', mask_real_copy)
        cv2.waitKey(0)
    print('mean of miou: ', miou_sum/len(img_names))
